/*
 * Mstore  engine
 *
 * IO engine using Mstore interface to test low-level performance of
 * Mstore.
 *
 */


#include <pthread.h>

#include "miniclient.h"



#include "config-host.h"
#include "fio.h"
#include "optgroup.h"
pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t cond = PTHREAD_COND_INITIALIZER;
pthread_mutex_t mutex2 = PTHREAD_MUTEX_INITIALIZER;
struct rpc_common
{
	struct rpc *mrpc;
	uint32_t sess_id; 
};
struct rpc_common *global_rpc = NULL;

struct mstore_data {
	struct rpc* mrpc;
	struct rpc_slot* io_ctx;
	struct io_u** aio_events;
	struct hash_table* file2blobid;
	uint32_t sess_id;
	int connected;
	int leader;
	struct flist_head completed_operations;
	uint64_t ops_scheduled;
	uint64_t ops_completed;
};

/*struct fio_rados_iou {
	struct flist_head list;
	struct thread_data* td;
	struct io_u* io_u;
	rados_completion_t completion;
	rados_write_op_t write_op;
};*/

struct fio_mstore_iou {
	struct flist_head list;
	struct thread_data* td;
	struct io_u* io_u;
};
char server_ip[20] = "10.0.0.37:12345";
struct rpc_session* common_sess;
struct rpc* common_rpc;
static void complete_callback(void* arg)
{
	struct fio_mstore_iou* fmi = (struct fio_mstore_iou*)arg;
	struct mstore_data* mstore = fmi->td->io_ops_data;

	//printf("callback\n");
	flist_add_tail(&fmi->list, &mstore->completed_operations);
	mstore->ops_completed++;

}

void
client_write_blob_complete(void* arg1, void* arg2 , void* arg3) 
{
	struct rpc* self = (struct rpc*) arg1;
	struct rpc_execution* ex = (struct rpc_execution*)arg2;
	struct msg_base *resp_h = get_response_header(ex);
	struct msg_base *req_h = get_request_header(ex);
	//decode user header
	struct dss_msg* dss_resp = (struct dss_msg*)resp_h->data;
	struct chunk_resp* ck_resp = (struct chunk_resp*)dss_resp->data;


	if (resp_h->status_code != SUCCESS) {
		printf("Failed to write blob %ld\n", ck_resp->id.blob_id);

	}
	else {
		//printf("Write blob %ld successfully\n", ms_h->blobid);
		complete_callback(arg3);
	}

	if(req_h->data_len < 1048576) free_msg_header(self, (void*)req_h);
	else free_large(self, req_h);
}

void client_read_blob_complete(void* arg1, void* arg2 , void* arg3) 
{
	struct rpc* self = (struct rpc*) arg1;
	struct rpc_execution* ex = (struct rpc_execution*)arg2;
	struct msg_base* resp_h = get_response_header(ex);
	struct msg_base* req_h = get_request_header(ex);
	
	//decode user header
	struct dss_msg* dss_resp = (struct dss_msg*)resp_h->data;
	struct chunk_resp* ck_resp = (struct chunk_resp*)dss_resp->data;


	if (resp_h->status_code != SUCCESS) {
		printf("Failed to write blob %ld\n", ck_resp->id.blob_id);
		
	}
	else {
		//printf("Read blob %ld successfully\n", ms_h->blobid);
		complete_callback(arg3);
	}
	//free resource, including header & slot..
	
	free_msg_header(self, req_h);


}
int cnt = 0;
void client_create_blob_complete(void* arg1, void* arg2 , void* arg3)
{
	struct rpc* self = (struct rpc*) arg1;
	struct rpc_execution* ex = (struct rpc_execution*)arg2;
	struct msg_base* req_h = get_request_header(ex);

	struct msg_base* resp_h = get_response_header(ex);
	//decode user header
	struct dss_msg* dss_resp = (struct dss_msg*)resp_h->data;
	struct chunk_resp* ck_resp = (struct chunk_resp*)dss_resp->data;

	uint64_t* blob_id = (uint64_t*)arg3;
	if (resp_h->status_code != SUCCESS) {
		*blob_id = -1;
	}
	else {
		*blob_id = ck_resp->id.blob_id;
		//printf("%d, Create new blob %ld successfully\n",cnt++, *blob_id);
	}

	//free resource, including header & slot..
	free_msg_header(self, req_h);


}

void client_delete_blob_complete(void* arg1, void* arg2 , void* arg3) {
	struct rpc* self = (struct rpc*) arg1;
	struct rpc_execution* ex = (struct rpc_execution*)arg2;
	struct msg_base* req_h = get_request_header(ex);

	struct msg_base* resp_h = get_response_header(ex);
	//decode user header
	struct dss_msg* dss_resp = (struct dss_msg*)resp_h->data;
	struct chunk_resp* ck_resp = (struct chunk_resp*)dss_resp->data;


	if (resp_h->status_code != SUCCESS) {
		printf("Failed to delete blob %ld\n", ck_resp->id.blob_id);
	}
	else {
		printf("Delete blob %ld successfully\n", ck_resp->id.blob_id);
	}
	//free resource, including header & slot..
	free_msg_header(self, req_h);

}
static int _fio_setup_mstore_data(struct thread_data* td, struct mstore_data**ms_data_ptr)
{
	struct mstore_data* mstore;

	if (td->io_ops_data)
		return 0;

	mstore = calloc(1, sizeof(struct mstore_data));
	if (!mstore)
		goto failed;

	mstore->connected = 0;
	void *mem = malloc(1 << 30);
	if(!mem) 
	{
		printf("Failed to malloc ..\n");
		exit(-1);
	}
	mstore->mrpc = rpc_create(mem, 1 << 30, RPC_RDMA | RPC_SHARED_RECV_QUEUE, NULL);
	mstore->file2blobid = create_hash_table();

	mstore->aio_events = calloc(td->o.iodepth, sizeof(struct io_u*));
	if (!mstore->aio_events)
		goto failed;
	
	INIT_FLIST_HEAD(&mstore->completed_operations);
	mstore->ops_scheduled = 0;
	mstore->ops_completed = 0;
	*ms_data_ptr = mstore;
	return 0;

failed:
	if (mstore) {
		if (mstore->aio_events)
			free(mstore->aio_events);
		free(mstore);
	}
	return 1;
}

static void _fio_mstore_rm_objects(struct thread_data* td, struct mstore_data* mstore)
{
	size_t i;
	for (i = 0; i < td->o.nr_files; i++) {
		struct fio_file* f = td->files[i];
		uint64_t blobid = (uint64_t)hash_find(mstore->file2blobid, (uint64_t)f->file_name);
		client_delete_blob(mstore->mrpc, mstore->sess_id, blobid, client_delete_blob_complete, NULL);
		while (run_poll_once(mstore->mrpc) <= 0);
	}
}

static void _session_close_accidently(struct rpc* self, void* usr) 
{	
	printf("goodbye \n");
}
static int _fio_mstore_connect(struct thread_data* td)
{
	struct mstore_data* mstore = td->io_ops_data;
	int r;
	const uint64_t file_size = 1u << 30;
	struct fio_file* f;
	uint32_t i;
	uint64_t blob_id;

	mstore->sess_id = rpc_create_session(mstore->mrpc, server_ip, NULL, NULL);
	if (mstore->sess_id < 0) {
	
		goto failed_early;
	}
	printf("nr_file = %d\n", td->o.nr_files);
	for (i = 0; i < td->o.nr_files; i++) {
		
		f = td->files[i];
		f->real_file_size = file_size;
		client_create_blob(mstore->mrpc, mstore->sess_id, client_create_blob_complete, &blob_id);
		while (run_poll_once(mstore->mrpc) <= 0);
		hash_insert(mstore->file2blobid, (uint64_t)f->file_name, (void*)blob_id);
	}
	return 0;

failed_obj_create:
	_fio_mstore_rm_objects(td, mstore);
	mstore->io_ctx = NULL;
failed_early:
	return 1;
}

static void _fio_mstore_disconnect(struct mstore_data* mstore)
{
	if (!mstore)
		return;
	if (mstore->sess_id) {
		rpc_close_session(mstore->mrpc, mstore->sess_id);
		return;
	}

}

static void fio_mstore_cleanup(struct thread_data* td)
{
	struct mstore_data* mstore = td->io_ops_data;
	if (mstore) {
		
		while (mstore->ops_scheduled != mstore->ops_completed) {
			run_poll_once(mstore->mrpc);
		}
			
		
		//_fio_mstore_rm_objects(td, mstore);
		//_fio_mstore_disconnect(mstore);
		free(mstore->aio_events);
		free(mstore);
	}
}


static enum fio_q_status fio_mstore_queue(struct thread_data* td,
	struct io_u* io_u)
{
	struct mstore_data* mstore = td->io_ops_data;
	struct fio_mstore_iou* fmi = io_u->engine_data;
	char* object = io_u->file->file_name;
	uint64_t blob_id = (uint64_t)hash_find(mstore->file2blobid, (uint64_t)object);

	
	fio_ro_check(td, io_u);
	
	if (io_u->ddir == DDIR_WRITE) {
		client_write_blob(mstore->mrpc, mstore->sess_id,blob_id, io_u->xfer_buf, io_u->offset, io_u->xfer_buflen, client_write_blob_complete, fmi);
		mstore->ops_scheduled++;
		return FIO_Q_QUEUED;
	}
	else if (io_u->ddir == DDIR_READ) {
		client_read_blob(mstore->mrpc,  mstore->sess_id,blob_id, io_u->xfer_buf, io_u->offset, io_u->xfer_buflen, client_read_blob_complete, fmi);
		mstore->ops_scheduled++;
		return FIO_Q_QUEUED;
	}
	else
	{
		log_err("WARNING: Only DDIR_READ, DDIR_WRITE are supported!");
		return FIO_Q_COMPLETED;
	}

}

static struct io_u* fio_mstore_event(struct thread_data* td, int event)
{
	struct mstore_data* mstore = td->io_ops_data;
	return mstore->aio_events[event];
}

int fio_mstore_getevents(struct thread_data* td, unsigned int min, unsigned int max, const struct timespec* t)
{
	struct mstore_data* mstore = td->io_ops_data;
	unsigned int events = 0;
	struct fio_mstore_iou* fmi;

	while (events < min) {
		while (flist_empty(&mstore->completed_operations)) {
			run_poll_once(mstore->mrpc);
		}
		//printf("completed op = %d\n", mstore->ops_completed);
		assert(!flist_empty(&mstore->completed_operations));

		fmi = flist_first_entry(&mstore->completed_operations, struct fio_mstore_iou, list);

		mstore->aio_events[events] = fmi->io_u;
		events++;
		flist_del(&fmi->list);
		if (events >= max) break;
	}
	return events;
}

static int fio_mstore_setup(struct thread_data* td)
{
	struct mstore_data* mstore = NULL;
	int r;

	/* allocate engine specific structure to deal with librados. */
	pthread_mutex_lock(&mutex);
	r = _fio_setup_mstore_data(td, &mstore);
	
	if (r) {
		log_err("fio_setup_mstore_data failed.\n");
		goto cleanup;
	}
	td->io_ops_data = mstore;

	/* Force single process mode.
	*/
	td->o.use_thread = 1;

	/* connect in the main thread to determine to determine
	* the size of the given RADOS block device. And disconnect
	* later on.
	*/
	r = _fio_mstore_connect(td);
	pthread_mutex_unlock(&mutex);
	if (r) {
		log_err("fio_mstore_connect failed.\n");
		goto cleanup;
	}
	mstore->connected = 1;

	return 0;
cleanup:
	fio_mstore_cleanup(td);
	return r;
}

/* open/invalidate are noops. we set the FIO_DISKLESSIO flag in ioengine_ops to
   prevent fio from creating the files
*/
static int fio_mstore_open(struct thread_data* td, struct fio_file* f)
{
	return 0;
}
static int fio_mstore_invalidate(struct thread_data* td, struct fio_file* f)
{
	return 0;
}

static void fio_mstore_io_u_free(struct thread_data* td, struct io_u* io_u)
{
	struct fio_mstore_iou* fmi = io_u->engine_data;

	if (fmi) {
		io_u->engine_data = NULL;
		fmi->td = NULL;

		free(fmi);
	}
}

static int fio_mstore_io_u_init(struct thread_data* td, struct io_u* io_u)
{
	struct fio_mstore_iou* fmi;
	fmi = calloc(1, sizeof(*fmi));
	fmi->io_u = io_u;
	fmi->td = td;
	INIT_FLIST_HEAD(&fmi->list);
	io_u->engine_data = fmi;
	return 0;
}


FIO_STATIC struct ioengine_ops ioengine = {
	.name = "mstore",
	.version = FIO_IOOPS_VERSION,
	.flags = FIO_DISKLESSIO,
	.setup = fio_mstore_setup,
	.queue = fio_mstore_queue,
	.getevents = fio_mstore_getevents,
	.event = fio_mstore_event,
	.cleanup = fio_mstore_cleanup,
	.open_file = fio_mstore_open,
	.invalidate = fio_mstore_invalidate,
	//.options = options,
	.io_u_init = fio_mstore_io_u_init,
	.io_u_free = fio_mstore_io_u_free,
	//.option_struct_size = sizeof(struct mstore_options),
};

static void fio_init fio_mstore_register(void)
{
	register_ioengine(&ioengine);
}

static void fio_exit fio_mstore_unregister(void)
{
	unregister_ioengine(&ioengine);
}